#ifndef RT_MEM_FLY_H
#define RT_MEM_FLY_H

#include <memory>
#include <vector>
#include <iostream>
#include "syncedmem.h"
#include "cuda_runtime.h"

namespace algocommon
{
    class Shape
    {
    public:
        Shape() : nums_(1), channels_(1), height_(1), width_(1), results_dim_(1){};
        Shape(const int nums, const int channels, const int height, const int width, const int results_dim)
            : nums_(nums), channels_(channels), height_(height), width_(width), results_dim_(results_dim){};

        // inline void Reshape(const int nums,const int channels,const int height,const int width)
        // {
        //     nums_=nums;
        //     channels_=channels;
        //     height_=height;
        //     width_=width;
        // }

        inline void Reshape(const int nums, const int channels, const int height, const int width, const int results_dim)
        {
            nums_ = nums;
            channels_ = channels;
            height_ = height;
            width_ = width;
            results_dim_ = results_dim;
        }

        inline int GetNums() const
        {
            return nums_;
        };

        inline int GetChannels() const
        {
            return channels_;
        };

        inline int GetHeight() const
        {
            return height_;
        };

        inline int GetWidth() const
        {
            return width_;
        };
        inline int GetResultsDim() const
        {
            return results_dim_;
        };

        friend std::ostream &operator<<(std::ostream &os, const Shape &shape)
        {
            os << "{"
               << "Num:" << shape.GetNums() << ","
               << "Channels:" << shape.GetChannels() << ","
               << "Height:" << shape.GetHeight() << ","
               << "Width:" << shape.GetWidth() << ","
               << "Results_dim:" << shape.GetResultsDim() << "}" << std::endl;
            return os;
        };

    private:
        int nums_;
        int channels_;
        int height_;
        int width_;
        int results_dim_;
    };
    //class TensorFly

    template <typename DataType>
    class TensorFly
    {
    public: //Tensor create
        TensorFly() : nums_(0), channels_(0), height_(0), width_(0), results_dim_(0),
                      one_image_count_(0), size_(0), total_count_(0), capacity_(0){};

        TensorFly(const int nums, const int channels, const int height, const int width)
            : nums_(nums), channels_(channels), height_(height), width_(width), capacity_(0)
        {
            one_image_count_ = channels_ * height_ * width_;
            total_count_ = nums_ * one_image_count_;
            size_ = sizeof(DataType) * total_count_;
            if (total_count_ > capacity_)
            {
                capacity_ = total_count_;
                data_.reset(new SyncedMemory(capacity_ * sizeof(DataType)));
            }
        };
        TensorFly(const int nums, const int channels, const int height, const int width, const int results_dim)
            : nums_(nums), channels_(channels), height_(height), width_(width), results_dim_(results_dim), capacity_(0)
        {
            one_image_count_ = channels_ * height_ * width_ * results_dim_;
            total_count_ = nums_ * one_image_count_;
            size_ = sizeof(DataType) * total_count_;
            if (total_count_ > capacity_)
            {
                capacity_ = total_count_;
                data_.reset(new SyncedMemory(capacity_ * sizeof(DataType)));
            }
        };
        TensorFly(const Shape &shape)
            : nums_(shape.GetNums()), channels_(shape.GetChannels()), height_(shape.GetHeight()),
              width_(shape.GetWidth()), results_dim_(shape.GetResultsDim()), capacity_(0)
        {
            one_image_count_ = channels_ * height_ * width_ * results_dim_;
            total_count_ = nums_ * one_image_count_;
            size_ = sizeof(DataType) * total_count_;
            if (total_count_ > capacity_)
            {
                capacity_ = total_count_;
                data_.reset(new SyncedMemory(capacity_ * sizeof(DataType)));
            }
        };

    public: //Tensor change
        void Resize(const int nums, const int channels, const int height, const int width)
        {
            nums_ = nums;
            channels_ = channels;
            height_ = height;
            width_ = width;

            one_image_count_ = channels_ * height_ * width_;
            total_count_ = nums_ * one_image_count_;
            size_ = sizeof(DataType) * total_count_;

            if (total_count_ > capacity_)
            {
                capacity_ = total_count_;
                data_.reset(new SyncedMemory(capacity_ * sizeof(DataType)));
            }
        }
        void Resize(const int nums, const int channels, const int height, const int width, const int results_dim)
        {
            nums_ = nums;
            channels_ = channels;
            height_ = height;
            width_ = width;
            results_dim_ = results_dim;
            one_image_count_ = channels_ * height_ * width_ * results_dim_;
            total_count_ = nums_ * one_image_count_;
            size_ = sizeof(DataType) * total_count_;

            if (total_count_ > capacity_)
            {
                capacity_ = total_count_;
                data_.reset(new SyncedMemory(capacity_ * sizeof(DataType)));
            }
        }
        void Reshape(const Shape &shape)
        {
            nums_ = shape.GetNums();
            channels_ = shape.GetChannels();
            height_ = shape.GetHeight();
            width_ = shape.GetWidth();
            results_dim_ = shape.GetResultsDim();

            one_image_count_ = channels_ * height_ * width_ * results_dim_;
            total_count_ = nums_ * one_image_count_;
            size_ = sizeof(DataType) * total_count_;

            if (total_count_ > capacity_)
            {
                capacity_ = total_count_;
                data_.reset(new SyncedMemory(capacity_ * sizeof(DataType)));
            }
        }

    public: //Tensor value
        inline size_t GetSize()
        {
            return size_;
        }
        inline int GetNums() const
        {
            return nums_;
        };

        inline int GetChannels() const
        {
            return channels_;
        };

        inline int GetHeight() const
        {
            return height_;
        };

        inline int GetWidth() const
        {
            return width_;
        };
        inline int GetResultsDim() const
        {
            return results_dim_;
        };

        inline std::string GetName() const
        {
            return name_;
        }

        inline void SetName(std::string &name)
        {
            name_ = name;
        }
        inline std::vector<int> &GetShapeList()
        {
            std::vector<int> ret{nums_, channels_, height_, width_};
            return std::move(ret);
        }
        inline Shape CreateShape()
        {
            return Shape(nums_, channels_, height_, width_);
        }

    public: //Tensor fly
        inline const DataType *cpu_data()
        {
            return reinterpret_cast<const DataType *>(data_->cpu_data());
        }

        inline const DataType *gpu_data()
        {
            return reinterpret_cast<const DataType *>(data_->gpu_data());
        }

        inline DataType *mutable_cpu_data()
        {
            return reinterpret_cast<DataType *>(data_->mutable_cpu_data());
        }

        inline DataType *mutable_cpu_data(int index)
        {
            return reinterpret_cast<DataType *>(data_->mutable_cpu_data()) + index * one_image_count_;
        }

        inline DataType *mutable_gpu_data()
        {
            return reinterpret_cast<DataType *>(data_->mutable_gpu_data());
        }
        inline DataType *mutable_gpu_data(int index)
        {
            return reinterpret_cast<DataType *>(data_->mutable_gpu_data()) + index * one_image_count_;
        }

        inline void memset_gpu_data()
        {
            return data_->memset_gpu_data();
        }

    private:
        int nums_;
        int channels_;
        int height_;
        int width_;
        int results_dim_;
        int one_image_count_;

        size_t size_;
        int total_count_;
        int capacity_;
        std::shared_ptr<SyncedMemory> data_;
        std::string name_;
    };

    typedef TensorFly<float> TensorFloat;
    typedef TensorFly<unsigned char> TensorUint8;
    typedef TensorFly<int> TensorInt32;

} // namespace algocommon
#endif